#! /usr/bin/python3
#
# memleak   Trace and display outstanding allocations to detect
#           memory leaks in user-mode processes and the kernel.
#
# USAGE: memleak [-h] [-p PID] [-t] [-a] [-o OLDER] [-c COMMAND]
#                [--combined-only] [--wa-missing-free] [-s SAMPLE_RATE]
#                [-T TOP] [-z MIN_SIZE] [-Z MAX_SIZE] [-O OBJ]
#                [interval] [count]
#
# Licensed under the Apache License, Version 2.0 (the "License")
# Copyright (C) 2016 Sasha Goldshtein.

from __future__ import annotations
from bcc import BPF
from time import sleep
from datetime import datetime
import resource
import argparse
import subprocess
import os
import sys
from typing import Dict, Tuple, Sequence, List, Any
import json
import traceback
from ctypes import *
import re
from shell_pipe import ShellPipe

p = ShellPipe()

class Allocation(object):
    def __init__(self, stack, size):
        self.stack = stack
        self.count = 1
        self.size = size

    def update(self, size):
        self.count += 1
        self.size += size

def run_command_get_output(command):
        p = subprocess.Popen(command.split(),
                stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        return iter(p.stdout.readline, b'')

def run_command_get_pid(command):
        p = subprocess.Popen(command.split())
        return p.pid

examples = """
EXAMPLES:

./memleak -p $(pidof allocs)
        Trace allocations and display a summary of "leaked" (outstanding)
        allocations every 5 seconds
./memleak -p $(pidof allocs) -t
        Trace allocations and display each individual allocator function call
./memleak -ap $(pidof allocs) 10
        Trace allocations and display allocated addresses, sizes, and stacks
        every 10 seconds for outstanding allocations
./memleak -c "./allocs"
        Run the specified command and trace its allocations
./memleak
        Trace allocations in kernel mode and display a summary of outstanding
        allocations every 5 seconds
./memleak -o 60000
        Trace allocations in kernel mode and display a summary of outstanding
        allocations that are at least one minute (60 seconds) old
./memleak -s 5
        Trace roughly every 5th allocation, to reduce overhead
"""

description = """
Trace outstanding memory allocations that weren't freed.
Supports both user-mode allocations made with libc functions and kernel-mode
allocations made with kmalloc/kmem_cache_alloc/get_free_pages and corresponding
memory release functions.
"""

parser = argparse.ArgumentParser(description=description,
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=examples)
parser.add_argument("-p", "--pid", type=int, default=-1,
        help="the PID to trace; if not specified, trace kernel allocs")
parser.add_argument("-t", "--trace", action="store_true",
        help="print trace messages for each alloc/free call")
parser.add_argument("interval", nargs="?", default=5, type=int,
        help="interval in seconds to print outstanding allocations")
parser.add_argument("count", nargs="?", type=int,
        help="number of times to print the report before exiting")
parser.add_argument("-a", "--show-allocs", default=False, action="store_true",
        help="show allocation addresses and sizes as well as call stacks")
parser.add_argument("-o", "--older", default=500, type=int,
        help="prune allocations younger than this age in milliseconds")
parser.add_argument("-c", "--command",
        help="execute and trace the specified command")
parser.add_argument("--combined-only", default=False, action="store_true",
        help="show combined allocation statistics only")
parser.add_argument("--wa-missing-free", default=False, action="store_true",
        help="Workaround to alleviate misjudgments when free is missing")
parser.add_argument("-s", "--sample-rate", default=1, type=int,
        help="sample every N-th allocation to decrease the overhead")
parser.add_argument("-T", "--top", type=int, default=10,
        help="display only this many top allocating stacks (by size)")
parser.add_argument("-z", "--min-size", type=int,
        help="capture only allocations larger than this size")
parser.add_argument("-Z", "--max-size", type=int,
        help="capture only allocations smaller than this size")
parser.add_argument("-O", "--obj", type=str, default="c",
        help="attach to allocator functions in the specified object")
parser.add_argument("--ebpf", action="store_true",
        help=argparse.SUPPRESS)
parser.add_argument("--percpu", default=False, action="store_true",
        help="trace percpu allocations")
parser.add_argument("--lib", default=False, action="store_true",
        help="collect by lib")
parser.add_argument("--func", default=False, action="store_true",
        help="collect by func")
parser.add_argument("-C", "--comm",
        help="trace the specified command")

args = parser.parse_args()

pid = args.pid
command = args.command
trace_all = args.trace
interval = args.interval
lib = args.lib
comm = args.comm
func = args.func

if command is not None:
        print("Executing '%s' and tracing the resulting process." % command)
        pid = run_command_get_pid(command)


bpf_source = ''

rootPath = os.path.dirname(os.path.abspath(__file__))
with open(os.path.join(rootPath, "mmap.c"), 'r') as f:
        bpf_source = f.read()
        bpf_source += '\n'

bpf_source = bpf_source.replace("SHOULD_PRINT", "1" if trace_all else "0")

stack_flags = "0"
#if not kernel_trace:
stack_flags += "|BPF_F_USER_STACK"
bpf_source = bpf_source.replace("STACK_FLAGS", stack_flags).replace("COMM_CARED", comm)

if args.ebpf:
    with open(os.path.join(rootPath, "mmap.out.c"), 'w') as f:
           f.write(bpf_source)

bpf = BPF(text=bpf_source)

def print_outstanding():
        print("[%s] Top %d stacks with outstanding allocations:" %
              (datetime.now().strftime("%H:%M:%S"), top_stacks))
        alloc_info = {}
        allocs = bpf["allocs"]
        stack_traces = bpf["stack_traces"]
        for address, info in sorted(allocs.items(), key=lambda a: a[1].size):
                # if BPF.monotonic_time() - min_age_ns < info.timestamp_ns:
                #         continue
                if info.stack_id < 0:
                        continue
                if info.stack_id in alloc_info:
                        alloc_info[info.stack_id].update(info.size)
                else:
                        stack = list(stack_traces.walk(info.stack_id))
                        combined = []
                        for addr in stack:
                                combined.append(bpf.sym(addr, pid,
                                        show_module=True, show_offset=True))
                        alloc_info[info.stack_id] = Allocation(combined,
                                                               info.size)
                if args.show_allocs:
                        print("\taddr = %x size = %s" %
                              (address.value, info.size))
        to_show = sorted(alloc_info.values(),
                         key=lambda a: a.size)[-top_stacks:]
        for alloc in to_show:
                print("\t%d bytes in %d allocations from stack\n\t\t%s" %
                      (alloc.size, alloc.count,
                       b"\n\t\t".join(alloc.stack).decode("ascii")))

def print_outstanding_combined():
        stack_traces = bpf["stack_traces"]
        stacks = sorted(bpf["combined_allocs"].items(),
                        key=lambda a: -a[1].total_size)
        entries = []
        for stack_id, info in stacks:
                taskcomm = info.comm.decode('ascii')
                if taskcomm != comm:
                        print('Error: task command cared fault: ', taskcomm)
                        continue
                try:
                        trace = []
                        for addr in stack_traces.walk(stack_id.value):
                                sym = bpf.sym(addr, pid,
                                                      show_module=True,
                                                      show_offset=True)
                                trace.append(sym)
                        trace = b"\n\t\t".join(trace).decode("ascii")
                except KeyError:
                        trace = "stack information lost"

                entry = ("\t%d bytes in %d allocations from stack\n\t\t%s" %
                         (info.total_size, info.number_of_allocs, trace))
                entries.append(entry)

        print("[%s] stacks with outstanding allocations:" %
              (datetime.now().strftime("%H:%M:%S")))

        print('\n'.join(reversed(entries)))


d3fileIndex = 0
def print_func_combined():
        global d3fileIndex
        d3fileIndex +=1

        allBytes = 0
        if c_uint32(pid) in bpf["stat"]:
                allBytes = bpf["stat"][c_uint32(pid)]
        print(f'{d3fileIndex} alloc size in bytes:{allBytes}')

        gf = open(os.path.join(rootPath, f"mmap.{d3fileIndex}gf.txt"), 'w')
      
        stack_traces = bpf["stack_traces"]
        stacks = sorted(bpf["combined_allocs"].items(),
                        key=lambda a: -a[1].total_size)

        for stack_id, info in stacks:
                taskcomm = info.comm.decode('ascii')
                if taskcomm != comm:
                        print('Error: task command cared fault: ', taskcomm)
                        continue
                try:
                        trace = []
                        for addr in stack_traces.walk(stack_id.value):
                                sym = bpf.sym(addr, pid,
                                        show_module=True,
                                        show_offset=True)
                                sym = str(sym)
                                trace.append(sym)
                        trace.reverse()
                        traces = ';'.join(trace)
                        gf.write(traces)
                        
                except KeyError:
                        print(f"stack information lost: {stack_id.value}")
                        traceback.print_exc()
                        gf.write(str(stack_id.value))
                gf.write(' ')
                gf.write(str(info.total_size))
                gf.write('\n')
        gf.close()
        print("end")

def print_lib_combined():
        global d3fileIndex
        d3fileIndex +=1

        allBytes = 0
        if c_uint32(pid) in bpf["stat"]:
                allBytes = bpf["stat"][c_uint32(pid)]
        print(f'{d3fileIndex} alloc size in bytes:{allBytes}')

        gf = open(os.path.join(rootPath, f"mmap.{d3fileIndex}lib.txt"), 'w')
      
        stack_traces = bpf["stack_traces"]
        stacks = sorted(bpf["combined_allocs"].items(),
                        key=lambda a: -a[1].total_size)

        for stack_id, info in stacks:
                taskcomm = info.comm.decode('ascii')
                if taskcomm != comm:
                        print('Error: task command cared fault: ', taskcomm)
                        continue
                try:
                        trace = []
                        for addr in stack_traces.walk(stack_id.value):
                                sym = bpf.sym(addr, pid,
                                        show_module=True,
                                        show_offset=True)
                                #print(f'sym {sym}')
                                s = sym.rfind(b'[')
                                e = sym.rfind(b']')
                                lib = sym[s+1:e].decode('ascii')
                                if lib == '':
                                        print(f'wann lib is emplty: {sym}')
                                if lib not in trace:
                                        trace.append(lib)
                        trace.reverse()
                        traces = ';'.join(trace)
                        gf.write(traces)
                        
                except KeyError:
                        print(f"stack information lost: {stack_id.value}")
                        traceback.print_exc()
                        gf.write(str(stack_id.value))
                gf.write(' ')
                gf.write(str(info.total_size))
                gf.write('\n')
        gf.close()
        print("end")

while True:
        if trace_all:
                print(bpf.trace_fields())
        else:
                try:
                        sleep(interval)
                except KeyboardInterrupt:
                        exit()
                p = f'pidof corerun' |p
                stdout,_ = p.communicate()
                if stdout != '':
                        pid = int(stdout)
                        print('wann', 'get pid', pid)
                else:
                        continue
                if lib:
                        print_lib_combined()
                elif func:
                        print_func_combined()
                else:
                        print_outstanding_combined()
                sys.stdout.flush()

